package top.luoyuanxiang.personalknowledgebase.service.impl;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;
import top.luoyuanxiang.personalknowledgebase.service.IDocumentEmbeddingService;

import java.util.List;

/**
 * 实现
 *
 * @author luoyuanxiang
 */
@Slf4j
@Service
public class DocumentEmbeddingServiceImpl implements IDocumentEmbeddingService {
    @Resource
    private VectorStore vectorStore;
    @Override
    public void addDocument(MultipartFile[] files) {
        // 读取文件，把文件分段然后放入向量库中
        for (MultipartFile file : files) {
            // 读取文件
            var pagePdfDocumentReader = new PagePdfDocumentReader(file.getResource());
            List<Document> documents = pagePdfDocumentReader.get();
            log.info("{} documents loaded", documents.size());

            // 2. split trunks
            List<Document> splitDocuments = TokenTextSplitter.builder()
                    .build()
                    .apply(documents);
            log.info("{} documents split", splitDocuments.size());

            // 3. create embedding and store to vector store
            log.info("create embedding and save to vector store");
            vectorStore.add(splitDocuments);

        }
    }

    @Override
    public List<Document> searchDocument(String query) {
        return vectorStore.similaritySearch(query);
    }
}
